// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/container/f14_map.h>
#include <memory>
#include <type_traits>

#include <pollux/common/base/simd_util.h>
#include <pollux/vector/lazy_vector.h>
#include <pollux/vector/simple_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    template<typename T>
    class DictionaryVector : public SimpleVector<T> {
    public:
        DictionaryVector(const DictionaryVector &) = delete;

        DictionaryVector &operator=(const DictionaryVector &) = delete;

        static constexpr bool can_simd = std::is_same_v<T, int64_t>;

        // Creates dictionary vector using base vector (dictionaryValues) and a set
        // of indices (dictionaryIndexArray).
        //
        // Base vector itself can be encoded as dictionary or constant. Base vector
        // can also be a lazy vector. If base vector is lazy and has not been
        // loaded yet, loading will be delayed until loaded_vector() is called.
        //
        // Base vector can contain duplicate values. This happens when dictionary
        // encoding is used to represent a result of a cardinality increasing
        // operator, for example, probe-side columns after cardinality increasing join
        // or result of an unnest.
        //
        // The number of indices can be less than the size of the base array, e.g. not
        // all elements of the base vector may be referenced. This happens when
        // dictionary encoding is used to represent a result of a filter or another
        // cardinality reducing operator, e.g. a selective join.
        DictionaryVector(
            pollux::memory::MemoryPool *pool,
            BufferPtr nulls,
            size_t length,
            VectorPtr dictionaryValues,
            BufferPtr dictionaryIndexArray,
            const SimpleVectorStats<T> &stats = {},
            std::optional<vector_size_t> distinctValueCount = std::nullopt,
            std::optional<vector_size_t> nullCount = std::nullopt,
            std::optional<bool> isSorted = std::nullopt,
            std::optional<ByteCount> represented_bytes = std::nullopt,
            std::optional<ByteCount> storageByteCount = std::nullopt);

        virtual ~DictionaryVector() override {
            dictionaryValues_->clear_containing_lazy_and_wrapped();
        }

        bool may_have_nulls() const override {
            POLLUX_DCHECK(initialized_);
            return BaseVector::nulls() || dictionaryValues_->may_have_nulls();
        }

        bool may_have_nulls_recursive() const override {
            POLLUX_DCHECK(initialized_);
            return BaseVector::may_have_nulls_recursive() ||
                   dictionaryValues_->may_have_nulls_recursive();
        }

        bool is_null_at(vector_size_t idx) const override;

        bool contains_null_at(vector_size_t idx) const override {
            if constexpr (std::is_same_v<T, ComplexType>) {
                if (is_null_at(idx)) {
                    return true;
                }

                auto innerIndex = getDictionaryIndex(idx);
                return dictionaryValues_->contains_null_at(innerIndex);
            } else {
                return is_null_at(idx);
            }
        }

        const T valueAtFast(vector_size_t idx) const;

        /**
         * @return the value at the given index value for a dictionary entry, i.e.
         * gets the dictionary value by its indexed value.
         */
        const T value_at(vector_size_t idx) const override {
            POLLUX_DCHECK(initialized_);
            POLLUX_DCHECK(!is_null_at(idx), "found null value at: {}", idx);
            auto innerIndex = getDictionaryIndex(idx);
            return scalarDictionaryValues_->value_at(innerIndex);
        }

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override;

        /**
         * Loads a SIMD vector of data at the virtual byteOffset given
         * Note this method is implemented on each vector type, but is intentionally
         * not virtual for performance reasons
         *
         * @param index at which to start the vector load
         * @return the vector of values starting at the given index
         */
        ksimd::batch<T> loadSIMDValueBufferAt(size_t index) const;

        inline const BufferPtr &indices() const {
            return indices_;
        }

        inline BufferPtr &indices() {
            return indices_;
        }

        const VectorPtr &value_vector() const override {
            return dictionaryValues_;
        }

        void set_value_vector(VectorPtr value_vector) override {
            setDictionaryValues(std::move(value_vector));
        }

        const BufferPtr &wrap_info() const override {
            return indices_;
        }

        void setWrapInfo(BufferPtr indices) override {
            indices_ = std::move(indices);
        }

        BufferPtr mutableIndices(vector_size_t size) {
            BaseVector::resize_indices(
                BaseVector::length_, size, BaseVector::pool_, indices_, &rawIndices_);
            return indices_;
        }

        uint64_t retained_size() const override {
            return BaseVector::retained_size() + dictionaryValues_->retained_size() +
                   indices_->capacity();
        }

        bool is_scalar() const override {
            return dictionaryValues_->is_scalar();
        }

        BaseVector *loaded_vector() override {
            if (initialized_) {
                return this;
            }

            SelectivityVector rows(dictionaryValues_->size(), false);
            for (vector_size_t i = 0; i < this->size(); i++) {
                if (!BaseVector::is_null_at(i)) {
                    auto ind = getDictionaryIndex(i);
                    rows.setValid(ind, true);
                }
            }
            rows.updateBounds();

            LazyVector::ensureLoadedRows(dictionaryValues_, rows);
            dictionaryValues_ = BaseVector::loaded_vector_shared(dictionaryValues_);
            setInternalState();
            return this;
        }

        const BaseVector *loaded_vector() const override {
            return const_cast<DictionaryVector<T> *>(this)->loaded_vector();
        }

        const BaseVector *wrapped_vector() const override {
            return dictionaryValues_->wrapped_vector();
        }

        vector_size_t wrapped_index(vector_size_t index) const override {
            return dictionaryValues_->wrapped_index(rawIndices_[index]);
        }

        std::string toString(vector_size_t index) const override {
            POLLUX_CHECK(
                initialized_,
                "Cannot convert to string because current DictionaryVector is not properly initialized yet.");
            if (BaseVector::is_null_at(index)) {
                return "null";
            }
            auto inner = rawIndices_[index];
            std::stringstream out;
            out << "[" << index << "->" << inner << "] "
                    << dictionaryValues_->toString(inner);
            return out.str();
        }

        void setDictionaryValues(VectorPtr dictionaryValues) {
            dictionaryValues_->clear_containing_lazy_and_wrapped();
            dictionaryValues_ = dictionaryValues;
            initialized_ = false;
            setInternalState();
        }

        /// Resizes the vector to be of size 'size'. If setNotNull is true
        /// the newly added elements point to the value at the 0th index.
        /// If setNotNull is false then the values and isNull is undefined.
        void resize(vector_size_t size, bool setNotNull = true) override {
            if (size > BaseVector::length_) {
                BaseVector::resize_indices(
                    BaseVector::length_,
                    size,
                    BaseVector::pool(),
                    indices_,
                    &rawIndices_);
            }

            // TODO Fix the case when base vector is empty.
            // https://github.com/facebookincubator/pollux/issues/7828

            BaseVector::resize(size, setNotNull);
        }

        VectorPtr slice(vector_size_t offset, vector_size_t length) const override;

        void validate(const VectorValidateOptions &options) const override;

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : BaseVector::pool_;
            return std::make_shared<DictionaryVector<T> >(
                selfPool,
                AlignedBuffer::copy(selfPool, BaseVector::nulls_),
                BaseVector::length_,
                dictionaryValues_->copy_preserve_encodings(),
                AlignedBuffer::copy(selfPool, indices_),
                SimpleVector<T>::stats_,
                BaseVector::distinctValueCount_,
                BaseVector::nullCount_,
                SimpleVector<T>::isSorted_,
                BaseVector::representedByteCount_,
                BaseVector::storageByteCount_);
        }

    private:
        // return the dictionary index for the specified vector index.
        inline vector_size_t getDictionaryIndex(vector_size_t idx) const {
            return rawIndices_[idx];
        }

        void setInternalState();

        BufferPtr indices_;
        const vector_size_t *rawIndices_ = nullptr;

        VectorPtr dictionaryValues_;
        // Caches 'dictionaryValues_.get()' if scalar type.
        SimpleVector<T> *scalarDictionaryValues_ = nullptr;
        // Caches 'scalarDictionaryValues_->getRawValues()' if 'dictionaryValues_'
        // is a FlatVector<T>.
        const T *rawDictionaryValues_ = nullptr;

        // Indicates whether internal state has been set. Can also be false if there
        // is an unloaded lazy vector under the encoding layers.
        bool initialized_{false};
    };

    template<typename T>
    using DictionaryVectorPtr = std::shared_ptr<DictionaryVector<T> >;
} // namespace kumo::pollux

#include <pollux/vector/dictionary_vector-inl.h>
